
# -*- encoding: utf-8 -*-
import redis
import ffmpeg
import os
from random import randint
import shutil
import numpy as np
from time import time
from hashlib import md5
from copy import copy
import cv2
from autoscript import pHash,cmpHash
from PIL import Image



def checkDir(srcdir,filetype='.mp4'):
    """输入文件佳验证

    Args:
        dir: 文件目录或者文件名，绝对路径
    return:
        srcfile: 视频文件列表，绝对路径
    """
    if not os.path.exists(srcdir):
        print("输入目录或输入文件不存在")
        exit(0)
     # 如果不是目录是单个文件
    if not os.path.isdir(srcdir):
        if not filetype in srcdir:
            print(f"不是{filetype}文件")
            exit(0)
        srcfiles = [srcdir,]
    # 如果是目录
    else:
        srcfiles = os.listdir(srcdir)
        # 找出所有mp4文件
        tempdir = srcfiles.copy()
        for file in tempdir:
            if not filetype in file:
                srcfiles.remove(file)
        # 没有mp4文件
        if len(srcfiles) == 0:
            print(f"该目录下无{filetype}文件")
            exit(0)
        # 返回视频文件列表
    return srcfiles


def selectVideo(pwd,selectnum = 20,cp=None):
    r = redis.Redis(host='127.0.0.1',port=6379,decode_responses=True)
    r.select(3)
    selector = []
    if r.exists(pwd):
        # 获取查重数组
        selector = r.smembers(pwd)
    filenum = len(os.listdir(pwd))
    # 数量少于选择的两倍则不选择
    if filenum < selectnum*2:
        print('数量过少，无法选择')
        return
    # 开始随机选择
    result = []
    while(True):
        # 生成长度为1×selectnum的随机数数组，随机范围1到filenum
        randindex = np.arange(1,filenum+1,1,dtype=np.int16)
        np.random.shuffle(randindex)
        randindex = randindex[:selectnum]
        # 重合数
        dupnum = 0
        # 检查路径
        print('检查路径')
        for index in randindex:
            # 拼接出选择视频的路径,看看是不是存在这个视频
            filepath = os.path.join(pwd,str(index)+'.mp4')
            if not os.path.exists(filepath):
                print(f'{filepath} 不存在')
                # 标记dupnum超出阈值，跳出重选
                dupnum = selectnum
                break
        # 开始查重
        print('开始查重')
        for oldselect in selector:
            # 只要历史中有一个重合数大于一半，跳出重选
            if dupnum > selectnum/2:
                break
            dupnum = 0
            for index in randindex:
                if str(index) in oldselect:
                    dupnum += 1
        # 跳出查重后dupnum仍然符合要求，则选取成功，跳出循环
        if dupnum <= selectnum/2:
            break
    # 将来本次选取写入数据库
    randindex = randindex.astype(str)
    r.sadd(pwd,'_'.join(randindex))
    # 如果复制目录位空，则复制到当前目录下
    if cp == None:
        cp = os.path.join(pwd,'_'.join(randindex))
        # 如果复制目录不存在
    if not os.path.exists(cp):
        os.mkdir(cp)
        print(f"路径 {cp} 不存在，已自动创建")
    # 开始复制
    for filename in randindex:
        filepath = os.path.join(pwd,filename + '.mp4')
        shutil.copy(filepath,cp)
        print(f"成功复制一个文件 {filepath} to {cp}")
    return randindex


# 获取视频图像序列
def getFrames(video_path):
    out, _ = (
    ffmpeg
    .input(video_path)
    .output('pipe:', format='rawvideo', pix_fmt='rgb24')
    .run(capture_stdout=True)
    )
    frames = (
        np
        .frombuffer(out, np.uint8)
        #.reshape([-1, height, width, 3])
    )
    return frames

# https://www.jianshu.com/p/b2bf18e71924
# 抽帧命令
# shuffleframes=0 1 2 3 4 -1 6 7 8 9 丢弃每10帧中的第5帧，改变index可交换帧
# -vsync vfr 改变时长模式，否则丢弃帧会用重复帧替换
# ffmpeg -i 1.mp4 -vf "shuffleframes=0 1 2 3 4 -1 6 7 8 9" -vsync vfr 2.mp4
# 计算帧数 ffprobe -v error -select_streams v:0 -count_packets -show_entries stream=nb_read_packets -of csv=p=0 2.mp4
def extractFrame(srcdir,dstdir,dropmode = None):
    srcfiles = checkDir(srcdir)
    if not os.path.exists(dstdir):
        os.mkdir(dstdir)
    if dropmode == None:
        # 每10帧丢弃第5帧
        dropmode = "0 1 2 3 4 -1 6 7 8 9"
    for inputvideo in list(srcfiles):
        print(f'正在对 {inputvideo} 进行抽帧')
        outputvideo = os.path.join(dstdir,inputvideo)
        inputvideo = os.path.join(srcdir,inputvideo)
        if not os.path.exists(inputvideo):
            print("输入的视频不存在")
            return
        out,_ = (
            ffmpeg.input(inputvideo)
            .filter("shuffleframes",dropmode)
            .output(outputvideo,vsync="vfr",vcodec="h264_nvenc",loglevel='quiet')
            .run()
        )
    
# 输入文件的文件 和 输出目录
# 参考输入文件名来分割
# 文件名为name-s_s_s.mp4，name为去重复原名字，s为分割点所在秒数
def trimVideo(dstdir, srcdir):
    if not os.path.exists(dstdir):
        if dstdir == None or dstdir == "":
            print("目标目录为空")
            return
        os.makedirs(dir)
    # 输入目录是否存在
    if not os.path.exists(srcdir):
        print("输入目录或输入视频不存在")
        return
    # 如果不是目录是单个文件
    if not os.path.isdir(srcdir):
        if not ".mp4" in srcdir:
            print("不是mp4文件")
            return
        srcfiles = [srcdir,]
    # 如果是目录
    else:
        srcfiles = os.listdir(srcdir)
        # 找出所有mp4文件
        tempdir = srcfiles.copy()
        for file in tempdir:
            if not ".mp4" in file:
                srcfiles.remove(file)
        # 没有mp4文件
        if len(srcfiles) == 0:
            print("该目录下无mp4文件")
            return
    # 开始处理
    for inputvideo in list(srcfiles):
        inputvideo = os.path.join(srcdir,inputvideo)
        if not os.path.exists(inputvideo):
            print("输入的视频不存在")
            return
        trimtime = inputvideo.split("\\")[-1].split(".")[0]
        #r如果没有分割名字，生成随机名字
        filename = str(time()).encode("utf-8")
        filename = md5(filename).hexdigest()[-10:] + "-"
        if "-" in trimtime:
            filename = trimtime.split("-")[0]+"-"
            trimtime = trimtime.split("-")[1]
        if "_" in trimtime:
            trimtime = trimtime.split("_")
        else:
            print("分割符不正确或无分割符")
            return
        trimtime.insert(0,"0")
        for i in range(len(trimtime)-1):
            outputvideo = os.path.join(dstdir, filename + trimtime[i]+ '_' + trimtime[i+1] + '.mp4')
            out,_ = (
            ffmpeg.input(inputvideo)
            # 是否启用gpu
            #.output(outputvideo, ss=trimtime[i], to=trimtime[i+1], avoid_negative_ts=1)
            .output(outputvideo, vcodec="h264_nvenc", ss=trimtime[i], to=trimtime[i+1], avoid_negative_ts=1)
            .run()
        )

# 镜像视频，delete是否删除源视频
def flipVideo(srcdir,delete = True):
    # 较验路径
    srcfiles = checkDir(srcdir)
        # 开始处理
    for inputvideo in list(srcfiles):
        outputvideo = os.path.join(srcdir,"copy-"+inputvideo)
        inputvideo = os.path.join(srcdir,inputvideo)
        if not os.path.exists(inputvideo):
            print("输入的视频不存在")
            return
        out,_ = (
            ffmpeg.input(inputvideo)
            # 是否启用gpu
            #.output(outputvideo, ss=trimtime[i], to=trimtime[i+1], avoid_negative_ts=1)
            .hflip()
            .output(outputvideo, vcodec="h264_nvenc")
            .run()
        )
        # stream = ffmpeg.input(inputvideo)
        # stream = ffmpeg.hflip(stream)
        # print(outputvideo)
        # stream.output(outputvideo).run()
        if delete:
            os.remove(inputvideo)
            os.rename(outputvideo,inputvideo)
 
 # 计算三通道感知哈希，三个值按GBR拼起来
def getGBRpHash(imagepath):
    if not os.path.exists(imagepath):
        print('图片目录不存在，无法计算感知哈希')
        return
    img = Image.open(imagepath)
    img = cv2.cvtColor(np.asarray(img),cv2.COLOR_RGB2BGR)
    img_g,img_b,img_r = cv2.split(img)
    img_gbr = [pHash(img_g), pHash(img_b), pHash(img_r)]
    return ''.join(img_gbr)
        
def checkDup(srcdir):
    """
    去重脚本
    1、连接redis
    2、提取关键帧，保存关键帧，名字为视频名字
    3、计算三通道感知哈希，按GBR拼接后,计算汉明距离
    4、计算出所有视频
    5、拿一个哈希值和和所有哈希值对比，相似的，关键帧全部放到一个文件夹
    """
    # 较验路径
    srcfiles = checkDir(srcdir)
    # 建一个保存i帧的文件夹
    framedir = os.path.join(srcdir,'Iframe')
    if os.path.exists(framedir):
        shutil.rmtree(framedir)
    os.mkdir(framedir)
    # 开始处理
    for inputvideo in list(srcfiles):
        outputimage = os.path.join(framedir,inputvideo[:-4] + '.jpg')
        inputvideo = os.path.join(srcdir,inputvideo)
        # keyarg = {
        #     'select_streams':'v'
        #     ,'show_entries':'frame=pkt_pts,pkt_dts,pkt_size,pkt_duration,pict_type'
        # }
        # # 查找关键帧位置，选择视频流、跳过不是关键帧、显示帧信息、只留pos位置偏移信息、格式化为json
        # # ffprobe -select_streams v -skip_frame nokey -show_frames -show_entries frame=pkt_pos -of json test.mp4
        # info = ffmpeg.probe(inputvideo,cmd='ffprobe',**keyarg)
        st = str(randint(2,4))
        out,_ = (
            ffmpeg.input(inputvideo,ss=f"00:00:0{st}") # 第2秒开始
            .filter('select','eq(pict_type,I)') # 选择i帧
            .output(outputimage, vframes=1, loglevel='quiet') # 提取1帧,不需要使用GPU
            .run()
        )
        # ffmpeg 命令 ffmpeg -ss 00:00:5 -i test.mp4 -vf select='eq(pict_type\,I)' -s 720*480 -vframes 1 iframe.jpg
        # -ss 从第5秒开始、-i 输入视频、-vf 过滤器选择I帧、-s resize、-vframes 导出1帧、输出文件名
        print(f'提取了I帧 {outputimage}')
    print('************** 开始计算hash *************')
    #连接redis,清空键值
    # rds = redis.Redis(host='127.0.0.1',port=6379,decode_responses=True)
    # rds.select(3)
    # rds.delete('iframes-images')
    # 获取所有的关键帧图片路径
    iframes = checkDir(framedir,filetype='.jpg')
    # 储存哈希值的列表
    hashlist = []
    for img in iframes:
        img_path = os.path.join(framedir,img)
        # 开始计算三通道哈希值
        hash = getGBRpHash(img_path)
        # rds.hset('iframes-images',img_path,hash)
        print(f'哈希计算完成 {img_path}')
        hashlist.append(hash)
    print('************** 开始查重 *************')
    # 迭代去除key 和 valuek
    # 考虑数量很多的情况所有加了redis,但好像多余了
    # img_hash = rds.hscan_iter('iframes-images')
    # templist = copy(hashlist)
    while len(hashlist) > 0:
        # 从列表末端开始遍历，每次取最后一个值，列表长度减1
        t_hash = hashlist.pop()
        # 提前创建重复收录文件夹,如果存在了就先删掉
        t_img = os.path.join(framedir,iframes.pop())
        t_dir = t_img[:-4]
        if os.path.exists(t_dir):
            shutil.rmtree(t_dir)
        os.mkdir(t_dir)
        # 重复标志，如果没有重复，就不需要这个文件夹，就要删掉
        isdup = False
        i = 0
        print(f'正在查询-----> {t_img}')
        # 这个循环的长度也会变化
        while i < len(hashlist):
            # 如果汉明距离相似率小于0.7则判定为不重复，跳过
            hanlen = cmpHash(t_hash,hashlist[i])
            if hanlen < 0.8:
                i += 1
                continue
            else:
                 # 发现重复,把它移到文件夹
                 dupimg = os.path.join(framedir,f'{hanlen}_{iframes[i]}')
                 os.rename(os.path.join(framedir,iframes[i]),dupimg)
                 print(f'!!!发现{t_img} 与 {dupimg} 重复!!!')
                 shutil.move(dupimg,t_dir)
                 isdup = True
                 # 文件列表和哈希列表中删除这个文件
                 iframes.pop(i)
                 hashlist.pop(i)
        # 循环结束后，如果有重复，把目标文件也移入文件夹,否则删除文件夹
        if isdup:
            shutil.move(t_img,t_dir)
        else:
            shutil.rmtree(t_dir)
            
# 查重后进行消重
def delDup(srcdir):
    # 较验路径
    srcfiles = checkDir(srcdir)
    # 如果Iframe 文件夹不存在，没有查重，停止
    iframepath = os.path.join(srcdir,'Iframe')
    if not os.path.exists(iframepath):
        print('请查重后再消重')
        exit(0)
    # 获取Iframe中的文件夹列表
    filelist = os.listdir(iframepath)
    dirlist = []
    for file in filelist:
        file = os.path.join(iframepath,file)
        if os.path.isdir(file):
            dirlist.append(file)
    if len(dirlist) == 0:
        print('Iframe 中没有找到重复文件夹')
    # 开始整理出要删除的视频的名字
    delnames = []
    for dirpath in dirlist:
        filelist = os.listdir(dirpath)
        for file in filelist:
            # 把带相似率的名字保存出来，把相似率和后缀去掉,去相似率时要注意文件名中原有_
            if '_' in file:
                file = '_'.join(file.split('_')[1:])
                delnames.append(file[:-4])
    if len(delnames) == 0:
        print('没有可删除的视频名字')
    # 开始删除
    for videofile in srcfiles:
        # 按理来说，视频和和甄名字只有后缀不同，所以取视频名字取后缀即可
        videoname = videofile[:-4]
        # 开始删除
        if videoname in delnames:
            print(f'!!! 删除视频 {videofile}')
            os.remove(os.path.join(srcdir,videofile))
                
    
    
def main():
    pwd = r"D:\Material\animal"
    cp =r"E:\File\Material\video\select3"
    srcdir = r'D:\Material\animal'
    dstdir = r'E:\File\Material\temp\select5'
    test = r'E:\File\Material\temp\srcdir'
    selectVideo(srcdir,20,cp)
main()
    